import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from pathlib import Path

def fig_ax_for_save(kwargs={}, name='', fontsize=20, height=10, width=10, ticks=True, axis_label=True):
    fig, ax= plt.subplots(subplot_kw=kwargs)
    if name:
        fig.suptitle(name, fontsize=fontsize)
        
    fig.set_figheight(height)
    fig.set_figwidth(width)
    if not axis_label:
        ax.set_yticklabels([])
        ax.set_xticklabels([])
    if not ticks:
        ax.set_xticks([])
        ax.set_yticks([])
    return fig, ax

def my_plot_force(f, save_path=None):
    fig, ax = fig_ax_for_save(ticks=False, axis_label=False)
    ax.set_aspect('equal', adjustable='box')
    im = ax.imshow(f[::-1])
    if not save_path is None: 
        fig.savefig(save_path, bbox_inches='tight')
    plt.close(fig)
    return 

def my_plot_surface(f, vmin, vmax, left=-1, right=1, save_path=None):
    n, m = f.shape
    xx, yy = np.meshgrid(np.linspace(left, right, m), np.linspace(left, right, n))
    fig, ax = fig_ax_for_save({"projection": "3d"})
    surf_pre = ax.plot_surface(xx, yy, f, cmap=cm.Spectral_r, vmin=vmin, vmax=vmax)
    cbar=plt.colorbar(surf_pre, shrink=0.85, ax=ax)
    cbar.ax.tick_params(labelsize=10)
    if not save_path is None: 
        fig.savefig(save_path, bbox_inches='tight')
    plt.close(fig)
    return 

def my_plot_ctf(f, levels, left=-1, right=1, save_path=None):
    n, m = f.shape
    xx, yy = np.meshgrid(np.linspace(left, right, m), np.linspace(left, right, n))
    
    fig, ax = fig_ax_for_save({}, ticks=False, axis_label=False)
    ax.set_aspect('equal', adjustable='box')
    ctf = ax.contourf(xx, yy, f, cmap=cm.Spectral_r, levels=levels)
    cbar=plt.colorbar(ctf, shrink=0.85, ax=ax, )
    cbar.ax.tick_params(labelsize=10)
    if not save_path is None: 
        fig.savefig(save_path, bbox_inches='tight')
    plt.close(fig)
    return 

def my_plot_ct(pre, ans, levels, left=-1, right=1, save_path=None):
    n, m = pre.shape
    xx, yy = np.meshgrid(np.linspace(left, right, m), np.linspace(left, right, n))
    
    fig, ax = fig_ax_for_save({}, ticks=False, axis_label=False)
    ax.set_aspect('equal', adjustable='box')
    
    ct1 = ax.contour(xx, yy, pre, colors='r', linestyles='dashed', linewidths=1.5, levels=levels)
    ct2 = ax.contour(xx, yy, ans, colors='b', linestyles='solid', linewidths=2, levels=levels)
    ax.clabel(ct1, inline=False, fontsize=20)
    ax.clabel(ct2, inline=False, fontsize=20)
    h1, _ = ct1.legend_elements()
    h2, _ = ct2.legend_elements()
    ax.legend([h1[0], h2[0]], ['Prediction', 'Reference'], prop={'size': 18})
    
    if not save_path is None: 
        fig.savefig(save_path, bbox_inches='tight')
    plt.close(fig)
    return 

def save_green_images(save_dir, F, pres, U, Diff, exp_name, diff_levels):
    cases_nums, n, _ = U.shape
    for i in range(cases_nums):
        save_path = Path(f"{save_dir}/{exp_name}/case{i+1}")
        if not save_path.is_dir():  
            save_path.mkdir(parents=True)
        
        f, pre, ans, diff = F[i], pres[i], U[i], Diff[i]
        my_plot_force(f, save_path/"force.png")
        my_plot_ctf(diff, levels=diff_levels, save_path=save_path/"ctf_diff.png")
        my_plot_ct(pre, ans, levels=np.linspace(ans.min(), ans.max(), 8)[1:-1], save_path=save_path/"ct.png")
        my_plot_surface(pre, vmin=ans.min(), vmax=ans.max(), save_path=save_path/"surface_pre.png")
        my_plot_surface(ans, vmin=ans.min(), vmax=ans.max(), save_path=save_path/"surface_ans.png")    

def save_FastSolver_images(save_dir, f, pre, ans):
    save_path = Path(save_dir)
    if not save_path.is_dir():  
        save_path.mkdir(parents=True)

    diff = np.abs(pre - ans)

    my_plot_force(f, save_path/"force.png")
    my_plot_ctf(diff, levels=50, save_path=save_path/"ctf_diff.png")
    my_plot_ctf(pre, levels=50, save_path=save_path/"ctf_pre.png")
    my_plot_ctf(ans, levels=50, save_path=save_path/"ctf_ans.png")
    my_plot_ct(pre, ans, levels=np.linspace(ans.min(), ans.max(), 8)[1:-1], save_path=save_path/"ct.png")    
    my_plot_surface(pre, vmin=ans.min(), vmax=ans.max(), save_path=save_path/"surface_pre.png")
    my_plot_surface(ans, vmin=ans.min(), vmax=ans.max(), save_path=save_path/"surface_ans.png")    


def draw_img(name, f, pre, ans, left, right):
    n, m = f.shape
    x = np.linspace(left, right, m)
    y = np.linspace(left, right, n)
    xx, yy = np.meshgrid(x, y)

    fig = plt.figure()
    fig.suptitle(name, fontsize=20)
    fig.set_figheight(20)
    fig.set_figwidth(30)

    ax1 = fig.add_subplot(2, 3, 1, aspect="equal")
    ax2 = fig.add_subplot(2, 3, 2, projection='3d')
    ax3 = fig.add_subplot(2, 3, 3, projection='3d')
    ax4 = fig.add_subplot(2, 3, 4, aspect="equal")
    ax5 = fig.add_subplot(2, 3, 5, aspect="equal")
    ax6 = fig.add_subplot(2, 3, 6, aspect="equal")

    im = ax1.imshow(f)
    ax1.set_title(f'$Source$', fontsize=20)
    cbar=plt.colorbar(im, shrink=0.85, ax=ax1)
    cbar.ax.tick_params(labelsize=10)

    ax2.set_title(f'$Prediction$', fontsize=20)
    surf_pre = ax2.plot_surface(xx, yy, pre, cmap=cm.Spectral_r,)
    cbar=plt.colorbar(surf_pre, shrink=0.85, ax=ax2)
    cbar.ax.tick_params(labelsize=10)

    ax3.set_title(f'$Reference$', fontsize=20)
    surf_ans = ax3.plot_surface(xx, yy, ans, cmap=cm.Spectral_r,)
    cbar=plt.colorbar(surf_ans, shrink=0.85, ax=ax3)
    cbar.ax.tick_params(labelsize=10)

    ax4.set_title(f'Difference', fontsize=20)
    ct = ax4.contourf(xx, yy, (ans - pre), cmap=cm.Spectral_r, levels=50)
    cbar=plt.colorbar(ct, shrink=0.85, ax=ax4)
    cbar.ax.tick_params(labelsize=10)

    ax5.set_title(f'$Prediction$', fontsize=20)
    ct1 = ax5.contourf(xx, yy, pre, alpha=1, cmap=cm.Spectral_r, levels=50)
    cbar=plt.colorbar(ct1, shrink=0.85, ax=ax5)
    cbar.ax.tick_params(labelsize=10)

    ax6.set_title(f'$Reference$', fontsize=20)
    ct2 = ax6.contourf(xx, yy, ans, alpha=1, cmap=cm.Spectral_r, levels=50)
    cbar=plt.colorbar(ct2, shrink=0.85, ax=ax6)
    cbar.ax.tick_params(labelsize=10)

    fig.tight_layout()
    plt.show()
    plt.close(fig)